import pickle
import pandas as pd
import argparse
import numpy as np
import itertools
parser = argparse.ArgumentParser()
parser.add_argument('--input_path', "-i", type=str,
                    default=r"G:\math_model_programm-main\result\K1_N2_lr0.0001_Problem_1.pkl")
parser.add_argument('--output_path', type=str, default=r'./output.xlsx')
args = parser.parse_args()
input_path = args.input_path
M, _, _, _ = pickle.load(open(input_path, 'rb'))
comlexity = M.comlexity
A_K = M._A_K
error = M.error
beta = M.beta

print(f"beta:{beta:.3f}")
print(f"error:{error:.3f}")
print(f"comlexity:{comlexity}")
dfs_real = [np.real(A_K_i) for A_K_i in A_K]
dfs_imag = [np.imag(A_K_i) for A_K_i in A_K]
for i, df_real, df_imag in zip(itertools.count(), dfs_real, dfs_imag):
    for l_real, l_imag in zip(df_real, df_imag):
        output = [f"{float(a):.3f}+{float(b):.3f}j" for a, b in zip(l_real, l_imag)]
        print("&".join(output)+r'\\')
    print(f"第{i}个A_K")
